#!/usr/bin/env python3
"""
快速内存转储方案 - 最简单的解决方案
"""

import subprocess
import os
import mmap
import glob

def step1_find_process():
    """步骤1: 找到Lingma进程"""
    print("🔍 步骤1: 查找Lingma进程")
    
    result = subprocess.run(['pgrep', '-f', 'Lingma'], capture_output=True, text=True)
    pid = result.stdout.strip()
    
    if not pid:
        print("❌ 未找到Lingma进程")
        print("请确保:")
        print("1. IntelliJ IDEA已启动")
        print("2. 通义灵码插件已激活")
        return None
    
    print(f"✅ 找到Lingma进程: PID = {pid}")
    
    # 检查进程信息
    try:
        cmd_result = subprocess.run(['ps', '-p', pid, '-o', 'comm='], 
                                  capture_output=True, text=True)
        print(f"   进程名: {cmd_result.stdout.strip()}")
    except:
        pass
    
    return pid

def step2_dump_memory(pid):
    """步骤2: 转储进程内存"""
    print(f"\n📥 步骤2: 转储PID {pid} 的内存")
    print("需要sudo权限...")
    
    # 使用gcore转储
    result = subprocess.run(['sudo', 'gcore', pid], 
                          capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"❌ 转储失败: {result.stderr}")
        return None
    
    # 查找生成的core文件
    core_file = f"core.{pid}"
    if not os.path.exists(core_file):
        # 尝试其他可能的位置
        core_files = glob.glob("core.*")
        if core_files:
            core_file = core_files[-1]
        else:
            print("❌ 未找到core文件")
            return None
    
    size_mb = os.path.getsize(core_file) / 1024 / 1024
    print(f"✅ 内存转储成功: {core_file} ({size_mb:.1f} MB)")
    return core_file

def step3_find_mapping_table(core_file):
    """步骤3: 在内存转储中查找映射表"""
    print(f"\n🔎 步骤3: 在 {core_file} 中搜索映射表")
    
    # 已知的映射
    known_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    found_tables = []
    
    print("扫描内存... (这可能需要几分钟)")
    
    with open(core_file, 'rb') as f:
        # 使用mmap提高性能
        with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as data:
            total_size = len(data)
            
            # 每100MB报告一次进度
            last_report = 0
            
            # 搜索256字节的段
            for offset in range(0, total_size - 256):
                # 进度报告
                if offset - last_report > 100 * 1024 * 1024:
                    progress = offset / total_size * 100
                    print(f"  进度: {progress:.1f}%")
                    last_report = offset
                
                # 8字节对齐可以加快搜索
                if offset % 8 != 0:
                    continue
                
                # 读取256字节
                try:
                    segment = data[offset:offset+256]
                    
                    # 快速检查 - 先检查一个已知映射
                    if segment[0xa1] != 0x6d:  # 如果'm'的映射不对，跳过
                        continue
                    
                    # 验证所有已知映射
                    matches = 0
                    for src, dst in known_mappings.items():
                        if segment[src] == dst:
                            matches += 1
                    
                    # 如果大部分映射匹配
                    if matches >= 6:
                        print(f"\n🎯 找到潜在映射表!")
                        print(f"   偏移: 0x{offset:x}")
                        print(f"   匹配: {matches}/8")
                        
                        found_tables.append((offset, matches, segment))
                        
                        # 如果完全匹配，可能就是我们要找的
                        if matches == 8:
                            print("   ✅ 完美匹配!")
                            return segment
                        
                except Exception as e:
                    # 忽略读取错误
                    pass
    
    print(f"\n找到 {len(found_tables)} 个候选映射表")
    
    # 返回最佳匹配
    if found_tables:
        found_tables.sort(key=lambda x: x[1], reverse=True)
        best_offset, best_matches, best_table = found_tables[0]
        
        print(f"\n使用最佳匹配 (偏移 0x{best_offset:x}, {best_matches}/8 匹配)")
        return best_table
    
    return None

def step4_decrypt(mapping_table):
    """步骤4: 使用映射表解密"""
    print("\n🔓 步骤4: 解密数据")
    
    if not os.path.exists('decoded_result.bin'):
        print("❌ 未找到 decoded_result.bin")
        print("请确保已运行Base64解码")
        return False
    
    # 保存映射表
    with open('found_mapping_table.bin', 'wb') as f:
        f.write(mapping_table)
    print("✅ 映射表已保存到 found_mapping_table.bin")
    
    # 读取加密数据
    with open('decoded_result.bin', 'rb') as f:
        encrypted = f.read()
    
    # 解密
    decrypted = bytearray()
    for byte in encrypted:
        decrypted.append(mapping_table[byte])
    
    # 检查结果
    # 转换为文本查看
    try:
        text = bytes(decrypted).decode('utf-8', errors='ignore')
        
        # 检查JSON特征
        if '{"' in text and '"questionText"' in text:
            print("✅ 解密成功!")
            
            # 保存结果
            with open('DECRYPTED_SUCCESS.json', 'wb') as f:
                f.write(decrypted)
            
            print("\n📄 解密结果已保存到: DECRYPTED_SUCCESS.json")
            
            # 显示预览
            print("\n预览 (前500字符):")
            print("-" * 60)
            print(text[:500])
            print("-" * 60)
            
            return True
        else:
            print("❓ 解密完成但内容可能不正确")
            
            # 仍然保存结果供检查
            with open('decrypted_attempt.bin', 'wb') as f:
                f.write(decrypted)
            
            print("结果已保存到 decrypted_attempt.bin")
            
            # 显示一些统计信息
            printable = sum(1 for b in decrypted if 32 <= b <= 126)
            print(f"\n统计:")
            print(f"  总字节数: {len(decrypted)}")
            print(f"  可打印字符: {printable} ({printable/len(decrypted)*100:.1f}%)")
            
            return False
            
    except Exception as e:
        print(f"❌ 解密出错: {e}")
        return False

def main():
    """主流程"""
    print("=" * 60)
    print("通义千问API解密 - 内存转储方案")
    print("=" * 60)
    
    # 步骤1: 找进程
    pid = step1_find_process()
    if not pid:
        return
    
    # 步骤2: 转储内存
    core_file = step2_dump_memory(pid)
    if not core_file:
        return
    
    # 步骤3: 查找映射表
    mapping_table = step3_find_mapping_table(core_file)
    if mapping_table is None:
        print("\n❌ 未找到映射表")
        print("\n可能的原因:")
        print("1. 映射表是动态生成的，还未生成")
        print("2. 映射表使用了不同的存储方式")
        print("3. 内存中的映射表已被清理")
        print("\n建议:")
        print("1. 在IDEA中触发一个API请求后立即运行此脚本")
        print("2. 尝试其他方案（如已知明文攻击）")
        return
    
    # 步骤4: 解密
    success = step4_decrypt(mapping_table)
    
    if success:
        print("\n🎉 恭喜! 解密成功完成!")
        print("\n下一步:")
        print("1. 查看 DECRYPTED_SUCCESS.json 获取完整内容")
        print("2. 使用 found_mapping_table.bin 解密其他请求")
    else:
        print("\n需要进一步分析")
    
    # 清理
    print(f"\n清理: 删除 {core_file}")
    try:
        os.remove(core_file)
    except:
        print(f"请手动删除: sudo rm {core_file}")

if __name__ == "__main__":
    main()